#!/usr/bin/env python3
"""
基于震动信号的虚拟键盘系统主程序

该程序整合了所有核心模块，实现了完整的虚拟键盘功能：
- 震动信号采集和处理
- 特征提取和键位识别
- 实时更新和模型适应
- 用户界面和交互
"""

import os
import sys
import yaml
import numpy as np
from typing import Dict, List, Optional, Tuple
from loguru import logger
import time
import threading
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from src.vibration_model import MultiModalVibration
from src.signal_processing import EndpointDetection, DTWAlgorithm, SignalPreprocessing
from src.feature_extraction import TimeDomainFeatures, FrequencyDomainFeatures
from src.key_recognition import KNNClassifier
from src.real_time_update import RealTimeUpdateSystem
from src.utils.config_manager import ConfigManager
from src.utils.data_manager import DataManager


class VibrationKeyboardSystem:
    """震动键盘系统主类"""
    
    def __init__(self, config_path: str = "config/config.yaml"):
        """
        初始化震动键盘系统
        
        Args:
            config_path: 配置文件路径
        """
        # 加载配置
        self.config_manager = ConfigManager(config_path)
        self.config = self.config_manager.get_config()
        
        # 设置日志
        self._setup_logging()
        
        # 初始化数据管理器
        self.data_manager = DataManager(self.config['data_storage'])
        
        # 初始化各个模块
        self._initialize_modules()
        
        # 系统状态
        self.is_running = False
        self.is_training = False
        
        logger.info("震动键盘系统初始化完成")
    
    def _setup_logging(self):
        """设置日志系统"""
        log_config = self.config['logging']
        
        # 移除默认处理器
        logger.remove()
        
        # 添加控制台处理器
        logger.add(
            sys.stdout,
            format=log_config['format'],
            level=log_config['level']
        )
        
        # 添加文件处理器
        log_path = Path(self.config['data_storage']['paths']['logs'])
        log_path.mkdir(parents=True, exist_ok=True)
        
        logger.add(
            log_path / "system.log",
            format=log_config['format'],
            level=log_config['level'],
            rotation=log_config['rotation'],
            retention=log_config['retention'],
            compression=log_config['compression']
        )
    
    def _initialize_modules(self):
        """初始化各个模块"""
        logger.info("开始初始化系统模块...")
        
        # 震动模型
        vibration_config = self.config['vibration_model']
        self.vibration_model = MultiModalVibration(
            horizontal_params=vibration_config['horizontal'],
            vertical_params=vibration_config['vertical'],
            combination_method="weighted_sum"
        )
        
        # 信号处理模块
        signal_config = self.config['signal_processing']
        self.endpoint_detector = EndpointDetection(
            energy_threshold_high=signal_config['endpoint_detection']['energy_threshold_high'],
            energy_threshold_low=signal_config['endpoint_detection']['energy_threshold_low'],
            zero_crossing_threshold=signal_config['endpoint_detection']['zero_crossing_threshold'],
            min_silence_length=signal_config['endpoint_detection']['min_silence_length'],
            min_speech_length=signal_config['endpoint_detection']['min_speech_length'],
            center_clip_threshold=signal_config['endpoint_detection']['center_clip_threshold'],
            frame_length=signal_config['energy_entropy']['frame_length'],
            frame_shift=signal_config['energy_entropy']['frame_shift'],
            sample_rate=self.config['hardware']['sensors']['accelerometer']['sample_rate']
        )
        
        self.dtw_algorithm = DTWAlgorithm(
            window_size=signal_config['dtw']['window_size'],
            step_pattern=signal_config['dtw']['step_pattern'],
            distance_metric=signal_config['dtw']['distance_metric'],
            normalization=signal_config['dtw']['normalization'],
            frame_length=signal_config['energy_entropy']['frame_length'],
            frame_shift=signal_config['energy_entropy']['frame_shift'],
            window_type=signal_config['energy_entropy']['window_type'],
            fft_size=signal_config['energy_entropy']['fft_size'],
            energy_constant=signal_config['energy_entropy']['energy_constant']
        )
        
        self.signal_preprocessor = SignalPreprocessing(
            sample_rate=self.config['hardware']['sensors']['accelerometer']['sample_rate'],
            normalize=True,
            denoise=True,
            filter_type="median"
        )
        
        # 特征提取模块
        feature_config = self.config['feature_extraction']
        self.time_features = TimeDomainFeatures(
            features=feature_config['time_domain']['features']
        )
        
        self.freq_features = FrequencyDomainFeatures(
            features=feature_config['frequency_domain']['features'],
            sample_rate=self.config['hardware']['sensors']['accelerometer']['sample_rate'],
            n_mfcc=feature_config['frequency_domain']['mfcc']['n_mfcc'],
            n_fft=feature_config['frequency_domain']['mfcc']['n_fft'],
            hop_length=feature_config['frequency_domain']['mfcc']['hop_length'],
            n_mels=feature_config['frequency_domain']['mfcc']['n_mels']
        )
        
        # 键位识别模块
        recognition_config = self.config['key_recognition']
        self.knn_classifier = KNNClassifier(
            n_neighbors=recognition_config['knn']['n_neighbors'],
            weights=recognition_config['knn']['weights'],
            algorithm=recognition_config['knn']['algorithm'],
            leaf_size=recognition_config['knn']['leaf_size'],
            metric=recognition_config['knn']['metric'],
            p=recognition_config['knn']['p']
        )
        
        # 实时更新模块
        update_config = self.config['real_time_update']
        self.real_time_updater = RealTimeUpdateSystem(
            max_samples_per_key=update_config['sample_update']['max_samples_per_key'],
            update_threshold=update_config['sample_update']['update_threshold'],
            outlier_detection=update_config['sample_update']['outlier_detection'],
            outlier_method=update_config['outlier_detection']['method'],
            contamination=update_config['outlier_detection']['contamination'],
            learning_rate=update_config['model_adaptation']['learning_rate'],
            adaptation_threshold=update_config['model_adaptation']['adaptation_threshold']
        )
        
        logger.info("所有模块初始化完成")
    
    def process_signal(self, signal_data: np.ndarray) -> Dict:
        """
        处理震动信号
        
        Args:
            signal_data: 原始震动信号
            
        Returns:
            处理结果字典
        """
        try:
            # 1. 信号预处理
            processed_signal = self.signal_preprocessor.preprocess(signal_data)
            
            # 2. 端点检测
            start_sample, end_sample = self.endpoint_detector.detect_endpoints(processed_signal)
            
            # 提取有效信号段
            if start_sample < end_sample:
                valid_signal = processed_signal[start_sample:end_sample]
            else:
                valid_signal = processed_signal
            
            # 3. 特征提取
            time_features = self.time_features.extract_features(valid_signal)
            freq_features = self.freq_features.extract_features(valid_signal)
            
            # 4. 特征组合
            combined_features = self._combine_features(time_features, freq_features)
            
            # 5. 键位识别
            if self.knn_classifier.is_trained:
                candidates = self.knn_classifier.predict_with_confidence(
                    combined_features.reshape(1, -1),
                    max_candidates=self.config['key_recognition']['classification']['max_candidates']
                )
                
                # 检查置信度阈值
                confidence_threshold = self.config['key_recognition']['classification']['confidence_threshold']
                filtered_candidates = [
                    (key, conf) for key, conf in candidates 
                    if conf >= confidence_threshold
                ]
                
                if not filtered_candidates:
                    filtered_candidates = candidates[:1]  # 至少返回一个候选
            else:
                filtered_candidates = [("unknown", 0.0)]
            
            result = {
                'signal_length': len(signal_data),
                'valid_length': len(valid_signal),
                'start_sample': start_sample,
                'end_sample': end_sample,
                'time_features': time_features,
                'freq_features': freq_features,
                'candidates': filtered_candidates,
                'processing_time': time.time()
            }
            
            return result
            
        except Exception as e:
            logger.error(f"信号处理失败: {e}")
            return {
                'error': str(e),
                'candidates': [("error", 0.0)]
            }
    
    def _combine_features(self, time_features: Dict, freq_features: Dict) -> np.ndarray:
        """
        组合时域和频域特征
        
        Args:
            time_features: 时域特征
            freq_features: 频域特征
            
        Returns:
            组合后的特征向量
        """
        # 提取数值特征
        time_values = list(time_features.values())
        
        # 处理频域特征（可能包含数组）
        freq_values = []
        for feature_name, feature_value in freq_features.items():
            if isinstance(feature_value, np.ndarray):
                # 对于数组特征，取均值或前几个值
                if feature_name == "power_spectral_density":
                    # 功率谱密度取前20个值
                    freq_values.extend(feature_value[:20])
                elif feature_name == "mfcc":
                    # MFCC取所有系数
                    freq_values.extend(feature_value)
                else:
                    # 其他特征取均值
                    freq_values.append(np.mean(feature_value))
            else:
                freq_values.append(feature_value)
        
        # 组合所有特征
        combined = np.array(time_values + freq_values)
        
        return combined
    
    def train_model(self, training_data: Dict[str, List[np.ndarray]]):
        """
        训练模型
        
        Args:
            training_data: 训练数据 {键位: [信号列表]}
        """
        logger.info("开始训练模型...")
        self.is_training = True
        
        try:
            # 准备训练数据
            X_train = []
            y_train = []
            feature_names = []
            
            for key, signals in training_data.items():
                for signal in signals:
                    # 处理信号
                    result = self.process_signal(signal)
                    
                    if 'error' not in result:
                        # 组合特征
                        combined_features = self._combine_features(
                            result['time_features'], 
                            result['freq_features']
                        )
                        
                        X_train.append(combined_features)
                        y_train.append(key)
            
            if len(X_train) > 0:
                X_train = np.array(X_train)
                y_train = np.array(y_train)
                
                # 生成特征名称
                time_feature_names = list(self.time_features.get_feature_names())
                freq_feature_names = []
                
                for feature_name in self.freq_features.get_feature_names():
                    if feature_name == "power_spectral_density":
                        freq_feature_names.extend([f"psd_{i}" for i in range(20)])
                    elif feature_name == "mfcc":
                        freq_feature_names.extend([f"mfcc_{i}" for i in range(self.config['feature_extraction']['frequency_domain']['mfcc']['n_mfcc'])])
                    else:
                        freq_feature_names.append(feature_name)
                
                feature_names = time_feature_names + freq_feature_names
                
                # 训练分类器
                self.knn_classifier.fit(X_train, y_train, feature_names)
                
                # 保存模型
                model_path = Path(self.config['data_storage']['paths']['models']) / "knn_model.pkl"
                model_path.parent.mkdir(parents=True, exist_ok=True)
                self.knn_classifier.save_model(str(model_path))
                
                logger.info(f"模型训练完成，样本数={len(X_train)}，类别数={len(set(y_train))}")
            else:
                logger.warning("没有有效的训练数据")
                
        except Exception as e:
            logger.error(f"模型训练失败: {e}")
        finally:
            self.is_training = False
    
    def load_model(self, model_path: str = None):
        """
        加载预训练模型
        
        Args:
            model_path: 模型文件路径
        """
        if model_path is None:
            model_path = Path(self.config['data_storage']['paths']['models']) / "knn_model.pkl"
        
        try:
            self.knn_classifier.load_model(str(model_path))
            logger.info(f"模型已从 {model_path} 加载")
        except Exception as e:
            logger.error(f"模型加载失败: {e}")
    
    def update_model(self, signal_data: np.ndarray, correct_key: str):
        """
        实时更新模型
        
        Args:
            signal_data: 震动信号
            correct_key: 正确的键位
        """
        try:
            # 处理信号
            result = self.process_signal(signal_data)
            
            if 'error' not in result:
                # 组合特征
                combined_features = self._combine_features(
                    result['time_features'], 
                    result['freq_features']
                )
                
                # 更新模型
                self.real_time_updater.update_sample(
                    self.knn_classifier,
                    combined_features.reshape(1, -1),
                    correct_key
                )
                
                logger.info(f"模型已更新，正确键位: {correct_key}")
            
        except Exception as e:
            logger.error(f"模型更新失败: {e}")
    
    def start_system(self):
        """启动系统"""
        if not self.is_running:
            self.is_running = True
            logger.info("震动键盘系统已启动")
            
            # 这里可以启动实时信号采集线程
            # self._start_signal_acquisition()
    
    def stop_system(self):
        """停止系统"""
        if self.is_running:
            self.is_running = False
            logger.info("震动键盘系统已停止")
    
    def get_system_status(self) -> Dict:
        """
        获取系统状态
        
        Returns:
            系统状态字典
        """
        status = {
            'is_running': self.is_running,
            'is_training': self.is_training,
            'model_trained': self.knn_classifier.is_trained,
            'available_keys': list(self.knn_classifier.classes_) if self.knn_classifier.classes_ is not None else [],
            'config': {
                'hardware': self.config['hardware'],
                'vibration_model': self.config['vibration_model'],
                'signal_processing': self.config['signal_processing'],
                'key_recognition': self.config['key_recognition']
            }
        }
        
        return status
    
    def simulate_key_press(self, key_position: Tuple[float, float], amplitude: float = 1.0) -> Dict:
        """
        模拟键位敲击
        
        Args:
            key_position: 键位坐标 (x, y)
            amplitude: 敲击振幅
            
        Returns:
            识别结果
        """
        try:
            # 模拟震动传播
            grid_size = (50, 50)
            grid_spacing = 0.01
            
            # 生成模拟信号（这里简化处理）
            sample_rate = self.config['hardware']['sensors']['accelerometer']['sample_rate']
            duration = 0.5  # 0.5秒
            num_samples = int(sample_rate * duration)
            
            # 生成震动信号（模拟）
            t = np.linspace(0, duration, num_samples)
            frequency = 50  # 50Hz震动
            signal_data = amplitude * np.sin(2 * np.pi * frequency * t) * np.exp(-5 * t)
            
            # 添加噪声
            noise = np.random.normal(0, 0.1, num_samples)
            signal_data += noise
            
            # 处理信号
            result = self.process_signal(signal_data)
            
            # 添加模拟信息
            result['simulation'] = {
                'key_position': key_position,
                'amplitude': amplitude,
                'frequency': frequency,
                'duration': duration
            }
            
            return result
            
        except Exception as e:
            logger.error(f"键位敲击模拟失败: {e}")
            return {'error': str(e)}
    
    def __str__(self) -> str:
        status = "运行中" if self.is_running else "已停止"
        trained = "已训练" if self.knn_classifier.is_trained else "未训练"
        return f"VibrationKeyboardSystem(status={status}, model={trained})"
    
    def __repr__(self) -> str:
        return self.__str__()


def main():
    """主函数"""
    # 创建系统实例
    system = VibrationKeyboardSystem()
    
    # 启动系统
    system.start_system()
    
    try:
        # 模拟一些测试
        print("系统状态:", system.get_system_status())
        
        # 模拟键位敲击
        print("\n模拟键位敲击测试:")
        for i in range(3):
            result = system.simulate_key_press((0.1 * i, 0.1 * i), amplitude=1.0)
            print(f"敲击 {i+1}: {result.get('candidates', [])}")
            time.sleep(1)
        
        print("\n系统运行完成")
        
    except KeyboardInterrupt:
        print("\n用户中断，正在停止系统...")
    finally:
        system.stop_system()


if __name__ == "__main__":
    main() 